#!/usr/bin/env python3
"""
动态线程池管理器
根据系统负载和任务特性动态调整线程池大小
"""

import os
import psutil
import time
import threading
from typing import List, Any, Callable, Optional
from concurrent.futures import ThreadPoolExecutor, Future
from dataclasses import dataclass
from queue import Queue, Empty


@dataclass
class ThreadPoolConfig:
    """线程池配置"""
    min_workers: int = 2
    max_workers: int = 16
    target_cpu_usage: float = 70.0  # 目标 CPU 使用率百分比
    target_memory_usage: float = 80.0  # 目标内存使用率百分比
    adjustment_interval: float = 5.0  # 调整间隔（秒）
    queue_size_threshold: int = 100  # 队列大小阈值


@dataclass
class PerformanceMetrics:
    """性能指标"""
    cpu_usage: float
    memory_usage: float
    active_threads: int
    queue_size: int
    throughput: float  # 每秒完成的任务数
    avg_task_time: float  # 平均任务时间


class DynamicThreadPoolManager:
    """动态线程池管理器"""
    
    def __init__(self, config: ThreadPoolConfig):
        """
        初始化动态线程池管理器
        
        Args:
            config: 线程池配置
        """
        self.config = config
        self.executor = ThreadPoolExecutor(max_workers=config.min_workers)
        self.task_queue = Queue()
        self.result_queue = Queue()
        
        # 性能监控
        self.metrics = PerformanceMetrics(0, 0, 0, 0, 0, 0)
        self.completed_tasks = 0
        self.total_task_time = 0
        self.start_time = time.time()
        
        # 控制线程
        self.monitor_thread = None
        self.adjuster_thread = None
        self.running = False
        
        # 锁
        self.lock = threading.Lock()
        
        # 任务统计
        self.task_stats = {
            'submitted': 0,
            'completed': 0,
            'failed': 0,
            'total_time': 0.0
        }
    
    def start(self):
        """启动线程池管理器"""
        self.running = True
        
        # 启动监控线程
        self.monitor_thread = threading.Thread(target=self._monitor_performance)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()
        
        # 启动调整线程
        self.adjuster_thread = threading.Thread(target=self._adjust_pool_size)
        self.adjuster_thread.daemon = True
        self.adjuster_thread.start()
    
    def stop(self):
        """停止线程池管理器"""
        self.running = False
        
        # 等待监控和调整线程结束
        if self.monitor_thread:
            self.monitor_thread.join(timeout=1.0)
        if self.adjuster_thread:
            self.adjuster_thread.join(timeout=1.0)
        
        # 关闭线程池
        self.executor.shutdown(wait=True)
    
    def submit_task(self, func: Callable, *args, **kwargs) -> Future:
        """
        提交任务到线程池
        
        Args:
            func: 任务函数
            *args: 位置参数
            **kwargs: 关键字参数
            
        Returns:
            Future: 任务未来对象
        """
        with self.lock:
            self.task_stats['submitted'] += 1
        
        # 包装任务以收集性能数据
        def wrapped_func(*args, **kwargs):
            start_time = time.time()
            try:
                result = func(*args, **kwargs)
                execution_time = time.time() - start_time
                
                with self.lock:
                    self.task_stats['completed'] += 1
                    self.task_stats['total_time'] += execution_time
                
                return result
            except Exception as e:
                execution_time = time.time() - start_time
                
                with self.lock:
                    self.task_stats['failed'] += 1
                    self.task_stats['total_time'] += execution_time
                
                raise e
        
        return self.executor.submit(wrapped_func, *args, **kwargs)
    
    def submit_batch(self, tasks: List[tuple]) -> List[Future]:
        """
        批量提交任务
        
        Args:
            tasks: 任务列表，每个元素是 (func, args, kwargs) 元组
            
        Returns:
            List[Future]: 未来对象列表
        """
        futures = []
        for func, args, kwargs in tasks:
            future = self.submit_task(func, *args, **kwargs)
            futures.append(future)
        return futures
    
    def _monitor_performance(self):
        """监控系统性能"""
        while self.running:
            try:
                # 获取系统指标
                cpu_percent = psutil.cpu_percent(interval=1.0)
                memory = psutil.virtual_memory()
                memory_percent = memory.percent
                
                # 获取线程池指标
                with self.lock:
                    active_threads = self.executor._max_workers
                    completed_tasks = self.task_stats['completed']
                    total_time = self.task_stats['total_time']
                
                # 计算吞吐量
                elapsed_time = time.time() - self.start_time
                throughput = completed_tasks / elapsed_time if elapsed_time > 0 else 0
                avg_task_time = total_time / completed_tasks if completed_tasks > 0 else 0
                
                # 更新指标
                self.metrics = PerformanceMetrics(
                    cpu_usage=cpu_percent,
                    memory_usage=memory_percent,
                    active_threads=active_threads,
                    queue_size=self.task_queue.qsize(),
                    throughput=throughput,
                    avg_task_time=avg_task_time
                )
                
                # 打印性能信息（可选）
                if self.config.verbose:
                    print(f"性能指标: CPU={cpu_percent:.1f}%, "
                          f"内存={memory_percent:.1f}%, "
                          f"线程={active_threads}, "
                          f"吞吐量={throughput:.2f}任务/秒")
                
                time.sleep(1.0)
                
            except Exception as e:
                print(f"性能监控错误: {e}")
                time.sleep(1.0)
    
    def _adjust_pool_size(self):
        """动态调整线程池大小"""
        while self.running:
            try:
                time.sleep(self.config.adjustment_interval)
                
                # 获取当前线程数
                current_workers = self.executor._max_workers
                
                # 计算新的线程数
                new_workers = self._calculate_optimal_workers()
                
                # 如果需要调整
                if new_workers != current_workers:
                    self._resize_thread_pool(new_workers)
                    if self.config.verbose:
                        print(f"调整线程池大小: {current_workers} -> {new_workers}")
                
            except Exception as e:
                print(f"线程池调整错误: {e}")
    
    def _calculate_optimal_workers(self) -> int:
        """
        计算最优线程数
        
        Returns:
            int: 最优线程数
        """
        current_workers = self.executor._max_workers
        
        # 基于 CPU 使用率调整
        if self.metrics.cpu_usage > self.config.target_cpu_usage:
            # CPU 使用率过高，减少线程
            new_workers = max(self.config.min_workers, 
                            int(current_workers * 0.8))
        elif self.metrics.cpu_usage < self.config.target_cpu_usage * 0.7:
            # CPU 使用率较低，可以增加线程
            new_workers = min(self.config.max_workers, 
                            int(current_workers * 1.2))
        else:
            new_workers = current_workers
        
        # 基于内存使用率调整
        if self.metrics.memory_usage > self.config.target_memory_usage:
            new_workers = max(self.config.min_workers, 
                            int(new_workers * 0.9))
        
        # 基于队列大小调整
        if (self.metrics.queue_size > self.config.queue_size_threshold and 
            self.metrics.cpu_usage < self.config.target_cpu_usage):
            new_workers = min(self.config.max_workers, new_workers + 2)
        
        return new_workers
    
    def _resize_thread_pool(self, new_size: int):
        """
        调整线程池大小
        
        Args:
            new_size: 新的线程池大小
        """
        # 创建新的线程池
        new_executor = ThreadPoolExecutor(max_workers=new_size)
        
        # 关闭旧线程池
        self.executor.shutdown(wait=False)
        
        # 替换执行器
        self.executor = new_executor
    
    def get_metrics(self) -> PerformanceMetrics:
        """
        获取性能指标
        
        Returns:
            PerformanceMetrics: 性能指标
        """
        return self.metrics
    
    def get_stats(self) -> dict:
        """
        获取统计信息
        
        Returns:
            dict: 统计信息
        """
        with self.lock:
            return self.task_stats.copy()
    
    def wait_completion(self, timeout: Optional[float] = None) -> bool:
        """
        等待所有任务完成
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            bool: 是否所有任务都已完成
        """
        start_time = time.time()
        
        while True:
            with self.lock:
                if (self.task_stats['submitted'] == 
                    self.task_stats['completed'] + self.task_stats['failed']):
                    return True
            
            if timeout and (time.time() - start_time) > timeout:
                return False
            
            time.sleep(0.1)


class BatchProcessor:
    """批处理器"""
    
    def __init__(self, 
                 pool_manager: DynamicThreadPoolManager,
                 batch_size: int = 10):
        """
        初始化批处理器
        
        Args:
            pool_manager: 线程池管理器
            batch_size: 批处理大小
        """
        self.pool_manager = pool_manager
        self.batch_size = batch_size
    
    def process_items(self, 
                      items: List[Any], 
                      process_func: Callable,
                      *func_args, 
                      **func_kwargs) -> List[Any]:
        """
        批量处理项目
        
        Args:
            items: 项目列表
            process_func: 处理函数
            *func_args: 函数参数
            **func_kwargs: 函数关键字参数
            
        Returns:
            List[Any]: 处理结果列表
        """
        results = []
        
        # 分批处理
        for i in range(0, len(items), self.batch_size):
            batch = items[i:i + self.batch_size]
            
            # 提交批处理任务
            futures = []
            for item in batch:
                future = self.pool_manager.submit_task(
                    process_func, item, *func_args, **func_kwargs
                )
                futures.append((item, future))
            
            # 收集结果
            for item, future in futures:
                try:
                    result = future.result()
                    results.append((item, result))
                except Exception as e:
                    results.append((item, e))
        
        return results


def test_dynamic_thread_pool():
    """测试动态线程池"""
    print("测试动态线程池...")
    
    # 创建配置
    config = ThreadPoolConfig(
        min_workers=2,
        max_workers=8,
        target_cpu_usage=70.0,
        adjustment_interval=2.0
    )
    config.verbose = True
    
    # 创建管理器
    manager = DynamicThreadPoolManager(config)
    manager.start()
    
    # 模拟任务
    def simulate_work(task_id: int, duration: float = 1.0):
        """模拟工作负载"""
        time.sleep(duration)
        return f"任务 {task_id} 完成"
    
    try:
        # 提交任务
        print("提交任务...")
        futures = []
        for i in range(20):
            future = manager.submit_task(simulate_work, i, 0.5)
            futures.append(future)
        
        # 等待完成
        print("等待任务完成...")
        for i, future in enumerate(futures, 1):
            result = future.result()
            print(f"  {i}. {result}")
        
        # 打印统计信息
        stats = manager.get_stats()
        print("\n任务统计:")
        for key, value in stats.items():
            print(f"  {key}: {value}")
        
        # 打印最终性能指标
        metrics = manager.get_metrics()
        print("\n最终性能指标:")
        print(f"  CPU 使用率: {metrics.cpu_usage:.1f}%")
        print(f"  内存使用率: {metrics.memory_usage:.1f}%")
        print(f"  活跃线程数: {metrics.active_threads}")
        print(f"  吞吐量: {metrics.throughput:.2f} 任务/秒")
        print(f"  平均任务时间: {metrics.avg_task_time:.2f} 秒")
        
    finally:
        manager.stop()


if __name__ == "__main__":
    test_dynamic_thread_pool()